import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from kmeans import KMeans

data = pd.read_csv('data/iris.csv') # 读取鸢尾花数据集

# 有几种不同类型的花
iris_types = ['SETOSA', 'VERSICOLOR', 'VIRGINICA']

#  x ,y 轴分别代表两个不同的特征
x_axis = 'petal_length'
y_axis = 'petal_width'
'''
plt.figure(figsize = (12,5))
plt.subplot(1,2,1)
for iris_type in iris_types:
    plt.scatter(data[x_axis][data['class']==iris_type],data[y_axis][data['class']==iris_type], label = iris_type)
plt.title('label_know')
plt.legend()
plt.subplot(1,2,2)
plt.scatter(data[x_axis][:], data[y_axis][:])
plt.title('label_unknown')
plt.show()
'''
num_examples = data.shape[0]
x_train = data[[x_axis,y_axis]].values.reshape(num_examples, 2)

# 指定训练所需要的参数
# k, 迭代次数
num_clusters = 3
max_iterations = 50
kmeans = KMeans(x_train, num_clusters)
centroids, closest_centroids_ids = kmeans.train(max_iterations)

plt.figure(figsize = (12,5))
plt.subplot(1,2,1)
for iris_type in iris_types:
    plt.scatter(data[x_axis][data['class']==iris_type],data[y_axis][data['class']==iris_type], label = iris_type)
plt.title('label_know')
plt.legend()
plt.subplot(1,2,2)
for centroid_id, centroid in enumerate(centroids):
    # 分类别进行遍历
    current_example_index = (closest_centroids_ids == centroid_id).flatten()
    plt.scatter(data[x_axis][current_example_index],data[y_axis][current_example_index], label = centroid_id)
for centroid_id, centroid in enumerate(centroids):
    plt.scatter(centroid[0],centroid[1],c = 'black', marker = 'x')
plt.title('kmeans')
plt.show()